import os
import numpy as np
from cv2 import cv2
from joblib import Parallel, delayed
from multiprocessing import cpu_count
import time

def safe_mkdir(subdir):
    if not os.path.exists(subdir):
        os.mkdir(subdir)

def get_dataset_subdirs(root_path):
    img_path = os.path.join(root_path, 'images')
    label_path = os.path.join(root_path, 'labels')
    return img_path, label_path

def read_img(f, mode='bgr'):
    if mode == 'bgr':
        return cv2.imread(f)
    else:
        return cv2.imread(f)[:, :, ::-1]

def read_label(f, img_size = None):
    H = W = 1
    if img_size is not None:
        H, W = img_size
    cls_ids, rois = [], []
    with open(f, 'r') as h:
        for line in h.readlines():
            cls_idx, ctrx, ctry, w, h = [float(v) for v in line.rstrip('\r\t\n ').split('\t')]
            x = int((ctrx - w / 2) * W)
            y = int((ctry - h / 2) * H)
            w = int(w * W)
            h = int(h * H)
            cls_ids.append(int(cls_idx))
            rois.append([x, y, w, h])
    return cls_ids, np.array(rois)

def create_color_map(class_num=5):
    cmap = [
        [255, 0,   0], # blue,    body
        [0, 255,   0], # green,   head
        [0,   0,   0], # black,   eye
        [0,   0, 255], # red,     ear
        [255, 0, 255], # magenta, tailroot
        [255, 255, 0],
        [0, 255, 255]
        ]
    return cmap[:class_num]

def xywh2vertices(rois):
    # rois: N_roi x 4, [[x1, y1, w, h]]
    x1, y1, w, h = rois[:, 0], rois[:, 1], rois[:, 2], rois[:, 3]
    # xi, yi: N_roi x 1
    x2, y2 = x1 + w, y1
    x3, y3 = x1, y1 + h
    x4, y4 = x1 + w, y1 + h
    # stack into N_roi x 2
    x1y1 = np.stack([x1, y1], axis=1)
    x2y2 = np.stack([x2, y2], axis=1)
    x3y3 = np.stack([x3, y3], axis=1)
    x4y4 = np.stack([x4, y4], axis=1)
    # stack into N-roi x 2 x 4
    vertices = []
    for pt1, pt2, pt3, pt4 in zip(x1y1, x2y2, x3y3, x4y4):
        vertices.append([list(pt1), list(pt2), list(pt3), list(pt4)])
    # reshape into N_roi x 4 x 2
    return np.array(vertices)

def vertices2xywh(rec):
    # rec: 4 x 2, four vertices of a box
    # Alert: it could be a rotated rectangle
    x1y1 = np.min(rec, axis=1)
    x2y2 = np.max(rec, axis=1)
    wh = x2y2 - x1y1
    # return xywh format, (type: list)
    xywh = []
    for pt, sz in zip(x1y1, wh):
        xywh.append([int(pt[0]), int(pt[1]), int(sz[0]), int(sz[1])])
    # return xywh format, (type: 2dlist, N x 4)
    return xywh

def warped_imgsize(M, src_imgsize):
    x, y = 0, 0
    w, h = src_imgsize
    rotater = np.mat(M[:, :2]).T
    bias = np.mat(M[:, 2])
    roi = np.array([[0, 0, w, h]])
    rec = xywh2vertices(roi)
    rec_rot = np.zeros_like(rec, dtype=float)
    vertices = np.mat(rec[0])
    rec_rot[0][:, :2] = vertices * rotater + bias
    roi_rot = vertices2xywh(rec_rot)
    offset_x, offset_y, w_new, h_new = roi_rot[0]
    return offset_x, offset_y, w_new, h_new

def rotatescaled_data(img, rois, angle, scale=1):
    H, W = img.shape[:2]
    # compile transform matrix:
    # [[alpha, -beta, (1-alpha) * x - beta * y]
    #  [-beta, alpha, beta * x - (a-alpha) * y]]
    M = cv2.getRotationMatrix2D((W/2, H/2), angle, scale)
    rotater = np.mat(M[:, :2]).T
    bias = np.mat(M[:, 2])
    img_rot = cv2.warpAffine(img, M, (W, H))
    # reshape rois into format: 
    # [[x1, y1]
    #  [x2, y2]
    #  [x3, y3]
    #  [x4, y4]] * n_roi
    rec_stack = xywh2vertices(rois) # ok
    rec_rot = np.zeros_like(rec_stack, dtype=float)
    for i, vertices in enumerate(rec_stack):
        vertices = np.mat(vertices)
        rec_rot[i][:, :2] = vertices * rotater + bias # ok
    rois_rot = vertices2xywh(rec_rot)
    return img_rot, rois_rot

def fliped_data(img, rois, flip_mode):
    H, W = img.shape[:2]
    img_flip = cv2.flip(img, flip_mode)
    rois_flip = np.copy(rois)
    if flip_mode == 0: # vertical flip
        rois_flip[:, 1] = H - rois[:, 1] - rois[:, 3]
    elif flip_mode == 1: # horizontal flip
        rois_flip[:, 0] = W - rois[:, 0] - rois[:, 2]
    elif flip_mode == -1: # both flip
        rois_flip[:, 0] = W - rois[:, 0] - rois[:, 2]
        rois_flip[:, 1] = H - rois[:, 1] - rois[:, 3]
    else:
        # print('[fliped_data.warning] invalid flip mode, return original data')
        img_flip = img
    return img_flip, rois_flip

def save_data(fp_img, fp_label, img_rot_flip, rois_rot_flip, cls_ids, W, H):
    cv2.imwrite(fp_img, img_rot_flip)
    with open(fp_label, 'w') as h_label:
        for cls_id, roi in zip(cls_ids, rois_rot_flip):
            h_label.write('%d' % cls_id)
            x, y, w, h = float(roi[0]), float(roi[1]), float(roi[2]), float(roi[3])
            cx = (x + w/2) / W
            cy = (y + h/2) / H
            w = w / W
            h = h / H
            h_label.write('\t%f\t%f\t%f\t%f\n' % (cx, cy, w, h))

def combined_conditions(angles, scales, flips):
    conditions = []
    for angle in angles:
        for scale in scales:
            for flip in flips:
                conditions.append((angle, scale, flip))
    return conditions

def show_rois(win, img, rois, cls_ids, cmap):
    for roi, cls_id in zip(rois, cls_ids):
        cv2.rectangle(img, rec=roi, color=cmap[cls_id])
    cv2.imshow(win, img)
    cv2.waitKey()

def save_augmented_data(img_outpath, label_outpath, f_img, f_label, img, rois, cls_ids, angle, scale, flip):
    fout_img = f_img.rstrip('.png') + '-r%d_s%.1f_f%d.png' % (angle, scale, flip)
    fpout_img = os.path.join(img_outpath, fout_img)
    fout_label = f_label.rstrip('.txt') + '-r%d_s%.1f_f%d.txt' % (angle, scale, flip)
    fpout_label = os.path.join(label_outpath, fout_label)
    if os.path.exists(fpout_img) and os.path.exists(fpout_label):
        print('condition [a%d]-[s%d]-[f%d] already exists' % (angle, scale, flip))
        return False
    img_rot, rois_rot = rotatescaled_data(img, rois, angle, scale=scale)
    rois_rot_ = np.array(rois_rot)
    img_rot_flip, rois_rot_flip = fliped_data(img_rot, rois_rot_, flip)
    # show_rois(win_roi, img_rot_flip, rois_rot_flip, cls_ids, cmap=cmap)
    save_data(fpout_img, fpout_label, img_rot_flip, rois_rot_flip, cls_ids, img.shape[1], img.shape[0])
    return True

def get_curr_time_str(tm=None):
    tm = time.time() if tm is None else tm
    return time.asctime(time.localtime(tm))

def data_augmentation(root_path, root_output):
    img_path, label_path = get_dataset_subdirs(root_path)
    img_outpath, label_outpath = get_dataset_subdirs(root_output)
    safe_mkdir(img_outpath)
    safe_mkdir(label_outpath)
    win_roi = 'preview'
    class_num = 5
    angles = [5, 10, 15, 20, -5, -10, -15, -20]
    scales = [0.6, 0.7, 0.8, 0.9]
    flips = [-2, 0, 1, -1] # h-flip, v-flip, hv-flips, original
    all_conditions = combined_conditions(angles, scales, flips)
    n_cpu = cpu_count()
    is_multi_thread = True
    print('augment condition number = ', len(all_conditions))

    cmap = create_color_map(class_num)
    start = time.time()
    print('start at ', get_curr_time_str(start))
    for f_img in os.listdir(img_path):
        if f_img.endswith('png'):
            fpath_img = os.path.join(img_path, f_img)
            f_label = f_img.rstrip('png') + 'txt'
            fpath_label = os.path.join(label_path, f_label)
            img = read_img(fpath_img)
            cls_ids, rois = read_label(fpath_label, img.shape[:2])
            if not is_multi_thread: # single-thread version
                [save_augmented_data(img_outpath, label_outpath, f_img, f_label, img, rois, cls_ids, angle, scale, flip) for angle, scale, flip in all_conditions]
            else: # multi-thread version
                Parallel(n_jobs=n_cpu)(delayed(save_augmented_data)(img_outpath, label_outpath, f_img, f_label, img, rois, cls_ids, angle, scale, flip) for angle, scale, flip in all_conditions)
    stop = time.time()
    print('stop at ', get_curr_time_str(stop))
    print('time cost = %.2f sec' % (stop - start))


if __name__ == '__main__':
    root_path = '/home/youngway/data/data_mouse'
    root_output = '/home/youngway/data/data_mouse_augmented'
    data_augmentation(root_path, root_output)
